# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import functools
import logging
import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.kr_ws.ws_parser_base import WsFeedHandlerBase
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.upbit_v1.kr_rest.product import UpbitProduct
from coin.exchange.upbit_v1.book_builder import BookBuilder


class UpbitFeedBase(WsFeedHandlerBase):
  def __init__(self):
    self.book = None

  def on_book(self, msg):
    self.book = msg

  def on_ticker(self, msg):
    pass

  def on_status(self, msg):
    pass

  def on_trade(self, msg):
    pass


TimestampObj = collections.namedtuple('TimestampObj', ['timestamp'])


@functools.lru_cache(maxsize=128)
def get_product_from_native_symbol(symbol):
  return UpbitProduct.FromNativeCrixStr(symbol)


def _parse_trade_side(side_str):
  if side_str == "BID":
    side = TradeSide.TRADE_BUY_SIDE
  elif side_str == "ASK":
    side = TradeSide.TRADE_SELL_SIDE
  else:
    raise ValueError('Unknown side: %s' % side_str)
  return side


def append_to_bids(bid_price, bid_qty, bids):
  if bid_price > 0 and bid_qty > 0:
    bids.append([bid_price, bid_qty])
  else:
    assert bid_price == 0 and bid_qty == 0


def append_to_asks(ask_price, ask_qty, asks):
  if ask_price > 0 and ask_qty > 0:
    asks.append([ask_price, ask_qty])
  else:
    assert ask_price == 0 and ask_qty == 0


def _is_book_valid(prev_book, curr_book):
  if prev_book is None:
    return True

  try:
    # If last_time_book is not complete, no need to compare, tolerate the case.
    prev_bid0 = prev_book.bid0().price
    prev_ask0 = prev_book.ask0().price
  except (AttributeError, IndexError):
    return True

  try:
    curr_bid0 = curr_book.bid0().price
    curr_ask0 = curr_book.ask0().price
  except (AttributeError, IndexError):
    return False

  spread_ratio = (curr_ask0 - curr_bid0) / curr_bid0
  if ((0.5 < curr_bid0 / prev_bid0 < 2) and (0.5 < curr_ask0 / prev_ask0 < 2)
      and spread_ratio < 0.05):
    return True
  return False


class UpbitHandler(WebsocketHandler):
  def on_ws_book(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    bids = []
    asks = []
    if 'type' in data:
      type = data['type']
      native_symbol = data['code']
    elif 'ty' in data:
      # New Upbit api use short style.
      type = data['ty']
      native_symbol = data['cd']
    else:
      raise ValueError(data)

    assert type == 'crixOrderbook', data
    product = get_product_from_native_symbol(native_symbol)
    native_symbol_crix = product.native_symbol_crix
    if not self.is_subscribed('book', native_symbol_crix):
      return

    if 'type' in data:
      for elem in data['orderbookUnits']:
        append_to_bids(elem['bidPrice'], elem['bidSize'], bids)
        append_to_asks(elem['askPrice'], elem['askSize'], asks)
    elif 'ty' in data:
      for elem in data['obu']:
        append_to_bids(elem['bp'], elem['bs'], bids)
        append_to_asks(elem['ap'], elem['as'], asks)
    else:
      raise ValueError(data)

    book = BookBuilder(native_symbol_crix)
    book.snapshot(bids=bids, asks=asks)
    book.update_timestamp(record.timestamp)

    prev_book = self._books.get(native_symbol_crix, None)
    if not _is_book_valid(prev_book, book):
      try:
        logging.error('%s book abnormal, bid0=%f, ask0=%f.' %
                      (native_symbol_crix, book.bid0().price, book.ask0().price))
      except Exception:
        logging.error('%s book abnormal and incomplete.' % native_symbol_crix)
      return

    self._books[native_symbol_crix] = book
    self.ready = True
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_ws_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())

    if 'type' in data:
      type = data['type']
      native_symbol = data['code']
      price = float(data['tradePrice'])
      qty = float(data['tradeVolume'])
      side = _parse_trade_side(data['askBid'])
    elif 'ty' in data:
      # New Upbit api use short style.
      type = data['ty']
      native_symbol = data['cd']
      price = float(data['tp'])
      qty = float(data['tv'])
      side = _parse_trade_side(data['ab'])
    else:
      raise ValueError(data)
    assert type == 'crixTrade', data
    product = get_product_from_native_symbol(native_symbol)
    native_symbol_crix = product.native_symbol_crix
    if not self.is_subscribed('trade', native_symbol_crix):
      return
    trade_event = create_market_trade_event(record.timestamp, native_symbol_crix, price, qty, side)
    if self._feed_checker:
      self._feed_checker.on_trade(product=product,
                                  trade_event=trade_event,
                                  timestamp=record.timestamp)
    self.publish_event('trade', native_symbol, trade_event)
